package com.tianji.aigc.service.impl;

import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.dashscope.app.Application;
import com.alibaba.dashscope.app.ApplicationParam;
import com.alibaba.dashscope.app.ApplicationResult;
import com.alibaba.dashscope.utils.JsonUtils;
import com.tianji.aigc.config.DashScopeProperties;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.ChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.TokenContext;
import com.tianji.common.utils.UserContext;
import io.reactivex.Flowable;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;


@Service
@RequiredArgsConstructor
@ConditionalOnProperty(prefix = "tj.ai", name = "chat-type", havingValue = "APP")
public class AppAgentChatServiceImpl implements ChatService {

    private final DashScopeProperties dashScopeProperties;


    private final ChatMemory chatMemory;
    private final ChatClient openAiChatClient;
    private final SystemPromptConfig systemPromptConfig;

//    private final JdbcChatMemory jdbcChatMemory;

    private final ChatSessionService chatSessionService;

    // 存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();

    //    输出结束的标记
    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();



    /*@Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
//        获取对话id
        var conversationId = ChatService.getConversationId(sessionId);
        String token = TokenContext.getToken();
        HashMap<String, Object> toolsMap = new HashMap<>();
        for (String tool : dashScopeProperties.getAppAgent().getTools()) {
            toolsMap.put(tool, MapUtil.of("user_token", token));
        }
        Map<String, Object> bizParams = MapUtil.<String, Object>builder()
                .put("user_defined_tokens", toolsMap)
                .build();

        ApplicationParam param = ApplicationParam.builder()
                .apiKey(dashScopeProperties.getKey())
                .appId(dashScopeProperties.getAppAgent().getId()) // 智能体id
                .prompt(question)
                .incrementalOutput(true) // 开启增量输出
                .bizParams(JsonUtils.toJsonObject(bizParams))
                .sessionId(conversationId) // 设置会话ID
                .build();

        Application application = new Application();
        try {
            Flowable<ApplicationResult> result = application.streamCall(param);

//            将Flowable 转化为 Flux 进行处理输出
            return Flux.from(result)
                    .doFirst(() -> { // 输出开始, 标记正在输出
                        GENERATE_STATUS.put(sessionId, true);
                    })
                    .doOnComplete(() -> { // 输出结束, 清除标记
                        GENERATE_STATUS.remove(sessionId);
                    })
                    .doOnError(throwable -> GENERATE_STATUS.remove(sessionId)) // 错误时清除标记
//                    输出过程中,判断是否正在输出, 如果正在输出, 则继续输出, 否则结束输出
                    .takeWhile(s -> Optional.ofNullable(GENERATE_STATUS.get(sessionId)).orElse(false))
                    .map(applicationResult -> {
//                        获取大模型的输出的内容
                        String text = applicationResult.getOutput().getText();
//                        封装响应对象
                        return ChatEventVO.builder()
                                .eventData(text)
                                .eventType(ChatEventTypeEnum.DATA.getValue())
                                .build();
                    })
                    .concatWith(Flux.just(STOP_EVENT));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }*/

    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
//         获取当前用户的token值
        var token = TokenContext.getToken();

//        获取对话id
        var conversationId = ChatService.getConversationId(sessionId);
//        大模型输出内容的缓存器, 用于在输出中断后的数据存储
        StringBuilder outputBuilder = new StringBuilder();

        // 获取用户id
        var userId = UserContext.getUser();

        var tools = this.dashScopeProperties.getAppAgent().getTools();
        var bizParams = Map.of("user_defined_tokens", CollStreamUtil.toMap(tools, tool -> tool, tool -> Map.of("user_token", token)));

//        设置APP请求参数
        ApplicationParam applicationParam = ApplicationParam.builder()
                .apiKey(this.dashScopeProperties.getKey())
                .appId(this.dashScopeProperties.getAppAgent().getId())
                .prompt(question)
                .bizParams(JsonUtils.toJsonObject(bizParams))
                .incrementalOutput(true) // 开启增强输出
                .sessionId(sessionId) // 用于会话记忆的
                .parameter("flow_stream_mode", "agent_format")
                .build();

//        创建APP对象并开始执行
        var application = new Application();

        try {
            Flowable<ApplicationResult> flowable = application.streamCall(applicationParam);

//            将Flowable转化为Flux进行流式输出
            return Flux.from(flowable)
                    .doFirst(() -> { // 输出开始标记正在输出
                        GENERATE_STATUS.put(sessionId, true);
                    })
                    .doOnComplete(() -> { // 输出结束清除标记
                        GENERATE_STATUS.remove(sessionId);
                    })
                    .doOnError(throwable -> GENERATE_STATUS.remove(sessionId)) // 错误时清除标记
                    .doOnCancel(() -> {
//                    当输出被取消时, 保存输出的内容到历史记录中
                        this.saveStopHistoryRecord(conversationId, outputBuilder.toString());
                    })
                    .doFinally(signalType -> {
//                    需要更新对话的标题 或 更新时间
                        var content = StrUtil.format("""
                             ------------
                            USER:{} \n
                            ASSISTANT:{}
                            ------------
                            """, question, outputBuilder.toString());
                        this.chatSessionService.update(sessionId, content, userId);
                    })
//                输出过程中, 判断是否正在输出, 如果正在输出, 则继续输出, 否则结束输出
                    .takeWhile(s -> Optional.ofNullable(GENERATE_STATUS.get(sessionId)).orElse(false))
                    .map(chatResponse -> {

//                    获取大模型的输出的内容
                        String text = chatResponse.getOutput().getText();
//                    追加到输出内容中
                        outputBuilder.append(text);
//                    封装响应对象
                        return ChatEventVO.builder()
                                .eventData(text)
                                .eventType(ChatEventTypeEnum.DATA.getValue())
                                .build();
                    })
                    .concatWith(Flux.just(STOP_EVENT));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void stop(String sessionId) {
//        移除标记
        GENERATE_STATUS.remove(sessionId);
    }



    /**
     * 保存输出停止的记录
     * @param conversationId      会话id
     * @param content             大模型输出内容
     */
    private void saveStopHistoryRecord(String conversationId, String content) {
        this.chatMemory.add(conversationId, new AssistantMessage(content));
//        this.jdbcChatMemory.add(conversationId, new AssistantMessage(content));
    }


    @Override
    public String chatText(String question) {
        return this.openAiChatClient.prompt()
                .system(promptSystem -> promptSystem.text(this.systemPromptConfig.getTextSystemMessage().get()))
                .user(question)
                .call()
                .content();
    }
}
